package com.air4.chinesetts.tts;

import android.content.Context;
import android.util.Log;

import com.air4.chinesetts.dispatcher.TtsStateDispatcher;
import com.air4.chinesetts.module.FastSpeech2;
import com.air4.chinesetts.module.MBMelGan;
import com.air4.chinesetts.utils.MD5Util;
import com.air4.chinesetts.utils.Processor;
import com.air4.chinesetts.utils.ThreadPoolManager;
import com.air4.chinesetts.utils.ZhProcessor;

import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author {@link "mailto:xuefeng.ding@outlook.com" "Xuefeng Ding"}
 * Created 2020-07-28 14:25
 */
class AudioWorker {
    private static final String TAG = "AudioWorker";

    private LinkedBlockingQueue<InputText> mInputQueue = new LinkedBlockingQueue<>();
    private InputText mCurrentInputText;
    private FastSpeech2 mFastSpeech2;
    private MBMelGan mMBMelGan;
    private Processor mProcessor;
    private TtsPlayer mTtsPlayer;
    private ZhProcessor zhProcessor;
    private Context context;

    AudioWorker(Context context, String fastspeech, String vocoder) {
        this.context = context;
        mFastSpeech2 = new FastSpeech2(fastspeech);
        mMBMelGan = new MBMelGan(vocoder);
        mProcessor = new Processor();
        mTtsPlayer = new TtsPlayer();
        zhProcessor = new ZhProcessor(context);

        ThreadPoolManager.getInstance().getSingleExecutor("worker").execute(() -> {
            //noinspection InfiniteLoopStatement
            while (true) {
                try {
                    mCurrentInputText = mInputQueue.take();
                    Log.d(TAG, "processing: " + mCurrentInputText.INPUT_TEXT);
                    TtsStateDispatcher.getInstance().onTtsStart(mCurrentInputText.INPUT_TEXT);
                    mCurrentInputText.proceed();
//                    TtsStateDispatcher.getInstance().onTtsStop();
                } catch (Exception e) {
                    Log.e(TAG, "Exception: ", e);
                }
            }
        });
    }

    void processInput(String inputText, float speed, boolean useCache) {
        Log.d(TAG, "add to queue: " + inputText);
        mInputQueue.offer(new InputText(inputText, speed, useCache));
    }

    void interrupt() {
        mInputQueue.clear();
        if (mCurrentInputText != null) {
            mCurrentInputText.interrupt();
        }
        mTtsPlayer.interrupt();
    }


    private class InputText {
        private final String INPUT_TEXT;
        private final float SPEED;
        private boolean isInterrupt;
        private boolean useCache;

        private InputText(String inputText, float speed, boolean useCache) {
            this.INPUT_TEXT = inputText;
            this.SPEED = speed;
            this.useCache = useCache;
        }

        private void proceed() {
            String[] sentences = INPUT_TEXT.split("[\n，。？?！!,;；]");
            Log.d(TAG, "speak: " + Arrays.toString(sentences));

            for (int i = 0; i < sentences.length; i++) {
                String sentence = sentences[i];

                long time = System.currentTimeMillis();

                float[] audioData;

                String fileName = MD5Util.MD5(sentence) + SPEED;
                File file = new File(TtsManager.getInstance().cacheDir, fileName);
                if (file.exists() && useCache) {
                    audioData = file2floatArray(file);
                } else {
//                    int[] inputIds = mProcessor.textToIds(sentence);
                    int[] inputIds = zhProcessor.text2ids(sentence);
                    TensorBuffer output = mFastSpeech2.getMelSpectrogram(inputIds, SPEED);

                    if (isInterrupt) {
                        Log.d(TAG, "proceed: interrupt");
                        return;
                    }
                    try {
                        audioData = mMBMelGan.getAudio(output);
                    } catch (Exception e) {
                        e.printStackTrace();
                        continue;
                    }
                    if (audioData != null && audioData.length > 0 && useCache) {
                        floatArray2File(audioData, file);
                    }
                }

                long encoderTime = System.currentTimeMillis();

                if (isInterrupt) {
                    Log.d(TAG, "proceed: interrupt");
                    return;
                }

                long vocoderTime = System.currentTimeMillis();

                Log.d(TAG, "Time cost: " + (encoderTime - time) + "+" + (vocoderTime - encoderTime) + "=" + (vocoderTime - time));

                mTtsPlayer.play(new TtsPlayer.AudioData(INPUT_TEXT, sentence, audioData, i == (sentences.length - 1)));
            }
        }

        private void interrupt() {
            this.isInterrupt = true;
        }
    }



    public static float[] file2floatArray(File tradeFile){
        try
        {
            FileInputStream fis = new FileInputStream(tradeFile);
            DataInputStream bos = new DataInputStream(fis);
            ArrayList<Float> floatArrayList = new ArrayList<>();
            while (bos.available() > 0)
            {
                floatArrayList.add(bos.readFloat());
            }
            fis.close();
            bos.close();

            float[] arr = new float[floatArrayList.size()];
            int index  = 0;
            for (Float v : floatArrayList) {
                arr[index++] = v;
            }
            return arr;
        }catch (FileNotFoundException e){
            e.printStackTrace();
        }catch (IOException e){
            e.printStackTrace();
        }
        return null;
    }


    public static boolean floatArray2File(float[] data, File file){
        OutputStream os = null;
        DataOutputStream dos = null;
        try {
            if (!file.exists()) {
                File parent = file.getParentFile();
                if (parent != null && !parent.exists()) {
                    if (!parent.mkdirs()) {
                        throw new IOException("File '" + file + "' could not be created");
                    }
                }
            }
            os = new FileOutputStream(file);
            dos = new DataOutputStream(os);

            for (float f: data) {
                dos.writeFloat(f);
            }
            dos.flush();

            return true;
        } catch (Exception e) {
            return false;
        } finally {
            try {
                if (dos != null) {
                    dos.flush();
                    dos.close();
                }
                if (os != null) {
                    os.close();
                }
            } catch (IOException e) {
            }
        }
    }
}